<?php
set_time_limit(180);
session_start();

header('Content-Type: application/json');

$action = $_GET['action'] ?? '';

switch ($action) {
    case 'key-status':
        echo json_encode(['isSet' => !empty($_SESSION['openrouter_api_key'])]);
        break;
    case 'modelscope-key-status':
        echo json_encode(['isSet' => !empty($_SESSION['modelscope_api_key'])]);
        break;
    case 'generate':
        handle_generate();
        break;
    default:
        http_response_code(404);
        echo json_encode(['error' => 'Action not found']);
}

function handle_generate() {
    $request_body = file_get_contents('php://input');
    $data = json_decode($request_body);

    if (json_last_error() !== JSON_ERROR_NONE) {
        http_response_code(400);
        echo json_encode(['error' => 'Invalid JSON body']);
        return;
    }

    $model = $data->model ?? '';
    $api_key = $data->apikey ?? '';
    
    // Determine which session key to set based on the model family
    if (!empty($api_key)) {
        if ($model === 'google/gemini-2.5-flash-image-preview') {
            $_SESSION['openrouter_api_key'] = $api_key;
        } else {
             // All other models in this UI use ModelScope
            $_SESSION['modelscope_api_key'] = $api_key;
        }
    }

    if ($model === 'google/gemini-2.5-flash-image-preview') {
        handle_openrouter_generation($data);
    } else {
        handle_modelscope_generation($data);
    }
}

function handle_openrouter_generation($data) {
    $api_key = $_SESSION['openrouter_api_key'] ?? '';
    if (empty($api_key)) {
        http_response_code(400);
        echo json_encode(['error' => 'OpenRouter API key not provided']);
        return;
    }

    $url = 'https://openrouter.ai/api/v1/chat/completions';

    $messages = [];
    $content = [['type' => 'text', 'text' => $data->prompt]];
    if (!empty($data->images)) {
        foreach ($data->images as $base64_image) {
            $content[] = ['type' => 'image_url', 'image_url' => ['url' => $base64_image]];
        }
    }
    $messages[] = ['role' => 'user', 'content' => $content];

    $post_data = [
        'model' => $data->model,
        'messages' => $messages,
    ];

    $ch = curl_init($url);
    curl_setopt($ch, CURLOPT_RETURNTRANSFER, true);
    curl_setopt($ch, CURLOPT_POST, true);
    curl_setopt($ch, CURLOPT_POSTFIELDS, json_encode($post_data));
    curl_setopt($ch, CURLOPT_HTTPHEADER, [
        'Content-Type: application/json',
        'Authorization: Bearer ' . $api_key,
        'HTTP-Referer: http://127.0.0.1', // It's a good practice to set a referer
        'X-Title: Gemini CLI Web UI',
    ]);

    $response = curl_exec($ch);
    $http_code = curl_getinfo($ch, CURLINFO_HTTP_CODE);
    $curl_error = curl_error($ch);
    curl_close($ch);

    if ($http_code >= 400 || $response === false) {
        $error_message = "OpenRouter API error: " . ($curl_error ? $curl_error : ($response ? $response : "HTTP code " . $http_code));
        http_response_code($http_code ?: 500);
        echo json_encode(['error' => $error_message]);
        return;
    }
    
    $response_data = json_decode($response, true);
    
    // Check for images in response (for models that might generate images)
    $images = $response_data['choices'][0]['message']['images'] ?? [];
    if (!empty($images) && isset($images[0]['image_url']['url'])) {
        echo json_encode(['imageUrl' => $images[0]['image_url']['url']]);
        return;
    }

    // Fallback to text response
    $text_response = $response_data['choices'][0]['message']['content'] ?? 'No valid image or text response found.';
    echo json_encode(['textResponse' => $text_response]);
}


function handle_modelscope_generation($data) {
    $api_key = $_SESSION['modelscope_api_key'] ?? '';
    if (empty($api_key)) {
        http_response_code(400);
        echo json_encode(['error' => 'ModelScope API key not provided']);
        return;
    }

    $base_url = 'https://api-inference.modelscope.cn/';
    $is_qwen_edit = $data->model === 'Qwen/Qwen-Image-Edit';
    $post_fields = [];

    // --- MODIFIED LOGIC START ---

    if ($is_qwen_edit) {
        // This model requires a flat structure with an image URL.
        if (empty($data->image_url)) {
            http_response_code(400);
            echo json_encode(['error' => 'Image URL is required for Qwen-Image-Edit.']);
            return;
        }
        $post_fields = [
            "model" => $data->model,
            "prompt" => $data->prompt,
            "image_url" => $data->image_url
        ];
    } else {
        // Other models get a nested 'parameters' object from the frontend.
        // We need to flatten it for the ModelScope API.
        $params = $data->parameters;
        $post_fields = [
            "model" => $data->model,
            "prompt" => $params->prompt,
            "negative_prompt" => $params->negative_prompt,
            "size" => $params->size,
            "steps" => $params->steps,
            "guidance_scale" => $params->guidance,
            "seed" => $params->seed
        ];
    }
    
    // --- MODIFIED LOGIC END ---


    // 1. Submit async task
    $ch = curl_init($base_url . "v1/images/generations");
    curl_setopt($ch, CURLOPT_POST, true);
    curl_setopt($ch, CURLOPT_POSTFIELDS, json_encode($post_fields, JSON_UNESCAPED_UNICODE));
    curl_setopt($ch, CURLOPT_HTTPHEADER, [
        "Content-Type: application/json",
        "Authorization: Bearer " . $api_key,
        "X-ModelScope-Async-Mode: true"
    ]);
    curl_setopt($ch, CURLOPT_RETURNTRANSFER, true);

    $response = curl_exec($ch);
    $http_code = curl_getinfo($ch, CURLINFO_HTTP_CODE);
    curl_close($ch);

    if ($http_code != 200) {
        http_response_code(500);
        echo json_encode(['error' => "Failed to submit task to ModelScope. HTTP {$http_code}: {$response}"]);
        return;
    }

    $response_data = json_decode($response, true);
    $task_id = $response_data["task_id"] ?? '';

    if (empty($task_id)) {
        http_response_code(500);
        echo json_encode(['error' => "Could not get task ID from ModelScope.", 'details' => $response_data]);
        return;
    }

    // 2. Poll for task result
    $max_attempts = 60; // Poll for up to 3 minutes (60 * 3s)
    $attempt = 0;
    while ($attempt < $max_attempts) {
        sleep(3);
        $ch = curl_init($base_url . "v1/tasks/{$task_id}");
        curl_setopt($ch, CURLOPT_HTTPHEADER, [
            "Authorization: Bearer " . $api_key,
            "X-ModelScope-Task-Type: image_generation"
        ]);
        curl_setopt($ch, CURLOPT_RETURNTRANSFER, true);
        $result = curl_exec($ch);
        $result_code = curl_getinfo($ch, CURLINFO_HTTP_CODE);
        curl_close($ch);

        if ($result_code == 200) {
            $task_data = json_decode($result, true);
            $task_status = $task_data["task_status"] ?? '';

            if ($task_status == "SUCCEED") {
                $image_url = $task_data["output_images"][0] ?? '';
                if (!empty($image_url)) {
                    echo json_encode(["imageUrl" => $image_url]);
                    return;
                }
            } elseif ($task_status == "FAILED") {
                http_response_code(500);
                echo json_encode(['error' => 'ModelScope task failed.', 'details' => $task_data]);
                return;
            }
        }
        $attempt++;
    }

    http_response_code(504); // Gateway Timeout
    echo json_encode(['error' => 'ModelScope task timed out after 3 minutes.']);
}
